/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: mc2 profiling
 * Author: huawei
 * Create: 2024-05-24
 */

#include "profiling_manager.h"
#include "toolchain/prof_api.h"
#include "toolchain/prof_common.h"
#include "log.h"
#include "sal_pub.h"
#include "common/sqe_context.h"
#include "common/aicpu_hccl_common.h"

namespace {
static constexpr u32 aging = 1;
}
namespace dfx {

std::mutex ProfilingManager::streamMutex_;
std::unordered_map<std::string, ProfCommInfo> ProfilingManager::tagOpInfoMap_;
std::unordered_map<s32, std::string> ProfilingManager::streamToTagMap_;

std::mutex ProfilingManager::startReportSqeIdxMutex_;
std::unordered_map<s32, u32> ProfilingManager::streamToSqeIdxMap_;

bool ProfilingManager::isL0Open_ = false;
bool ProfilingManager::isL1Open_ = false;

constexpr std::uint32_t HCCLINFO_REPORT_BATCH_NUM = 2;

bool ProfilingManager::IsProfOn(uint64_t feature)
{
    if (AdprofCheckFeatureIsOn == nullptr) {
        return false;
    }
    return AdprofCheckFeatureIsOn(feature) > 0;
}

bool ProfilingManager::IsL1fromOffToOn()
{
    if (((!GetProfL1State()) && ProfilingManager::IsProfL1On())) {
        HCCL_INFO("Profiling L1 switch form off to on.");
        return true;
    }
    return false;
}

bool ProfilingManager::IsProfL1On()
{
    if (IsProfOn(ADPROF_TASK_TIME_L1)) {
        isL1Open_ = true;
        return true;
    }
    isL1Open_ = false;
    return false;
}

bool ProfilingManager::IsProfL0On()
{
    if (IsProfOn(ADPROF_TASK_TIME_L0)) {
        isL0Open_ = true;
        return true;
    }
    isL0Open_ = false;
    return false;
}

bool ProfilingManager::GetProfL0State()
{
    if (!isL0Open_) {
        HCCL_DEBUG("[GetProfL0State] L0 is not open");
        return false;
    }
    return true;
}

bool ProfilingManager::GetProfL1State()
{
    if (!isL1Open_) {
        HCCL_DEBUG("[GetProfL1State] L1 is not open");
        return false;
    }
    return true;
}

HcclResult ProfilingManager::CallMsprofReportAdditionInfo(uint32_t type, uint64_t timeStamp, const void *data, int len)
{
    MsprofAdditionalInfo reporterData{};
    reporterData.level = MSPROF_REPORT_AICPU_LEVEL;
reporterData.type = type;
    reporterData.threadId = SalGetTid();
    reporterData.dataLen = len;
    reporterData.timeStamp = timeStamp;
    s32 sret = memcpy_s(reporterData.data, sizeof(reporterData.data), data, len);
    CHK_PRT_RET(sret != EOK, HCCL_ERROR("memcpy failed. errorno[%d]:", sret), HCCL_E_MEMORY);
    HCCL_DEBUG("CallMsprofReportAdditionInfo, AdditionInfoType[%u]", type);
    CHK_PTR_NULL(AdprofReportAdditionalInfo);
    CHK_PRT_RET(AdprofReportAdditionalInfo(aging, &reporterData, sizeof(MsprofAdditionalInfo)) != 0,
        HCCL_ERROR("AdprofReportAdditionalInfo failed."),
        HCCL_E_INTERNAL);
    HCCL_DEBUG("CallMsprofReportAdditionInfo with additionInfoType[%u] successfully", type);
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::ReportTaskInfo()
{
    SqeContext *context = GetSqeContext();
    auto ctx = AicpuGetComContext();
    CHK_PTR_NULL(context->buffPtr);
    MsprofAicpuHcclTaskInfo taskInfos[HCCLINFO_REPORT_BATCH_NUM] = {0};
    for (uint32_t streamId = 0; streamId < AC_MAX_RANK_NUM; streamId++) {
        auto &buff = context->buffPtr[streamId];
        uint16_t &lastSqeIdx = ctx->profilingExtendInfo.lastSqeIdxs[streamId];
        HCCL_INFO("Rank %u stream %u has %u sqes, sqeCnt: %u, lastSqeIdx: %u", ctx->rankId, streamId,
                  buff.tailSqeIdx, buff.sqeCnt, lastSqeIdx);
        auto endIdx = static_cast<uint32_t>(buff.tailSqeIdx);
        for (uint32_t idx = lastSqeIdx, batchId = 0; idx < endIdx; ++idx) {
            auto& taskInfo = taskInfos[batchId++];
            dfx::ProfilingExtendInfoHelper::Ctx2MsprofAicpuMC2HcclInfo(ctx, taskInfo);
            taskInfo.planeID = streamId;
            SqeInfo sqeInfo;
            SqeContextUtils::QuerySqeInfo(
                buff.localBuff + idx * AC_SQE_SIZE, buff.sqeType[idx], buff.addInfo[idx], &sqeInfo);
            dfx::ProfilingExtendInfoHelper::SqeInfo2MsprofAicpuMC2HcclInfo(sqeInfo, taskInfo);
            taskInfo.timeStamp = GetCurCpuTimestamp(true);
            DumpHcclInfo(taskInfo, batchId, idx);
            if (batchId == HCCLINFO_REPORT_BATCH_NUM || idx == (endIdx - 1)) {
                CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(MSPROF_REPORT_AICPU_MC2_BATCH_HCCL_INFO,
                    0, taskInfos, sizeof(MsprofAicpuHcclTaskInfo) * batchId));
                batchId = 0;
                memset_s(taskInfos, sizeof(taskInfos), 0, sizeof(taskInfos));
            }
        }
        lastSqeIdx = buff.tailSqeIdx;
    }
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::ReportTaskExecTimeLine(AicpuComProf *acprof)
{
    if (!GetProfL1State()) {
        return HCCL_SUCCESS;
    }
    CHK_PTR_NULL(aicpu::GetTaskAndStreamId);
    uint64_t taskId = 0U;
    uint32_t streamId = 0;
    CHK_PRT_RET(aicpu::GetTaskAndStreamId(taskId, streamId) != aicpu::status_t::AICPU_ERROR_NONE,
        HCCL_ERROR("Failed to get task id and stream id."),
        HCCL_E_PARA);

    // 正常一个kfc算子的展开任务不会超过`AC_MAX_PROF_COMM_CNT`轮
    u32 workRcdCnt = acprof->workCnt;
    CHK_PRT_RET(workRcdCnt > AC_MAX_PROF_COMM_CNT,
        HCCL_ERROR("WorkRcdCnt %u should not bigger than %u,", workRcdCnt, AC_MAX_PROF_COMM_CNT),
        HCCL_E_PARA);

    HCCL_INFO("workRcdCnt now is %u g_proxLoopCnt is %u", workRcdCnt, g_proxLoopCnt);
    for (u32 j = 0; j < workRcdCnt; j++) {
        AicpuComProfCommLoop *commRcd = &acprof->commLoop[j];
        AicpuKfcProfCommTurn commTurnProf;
        commTurnProf.serverStartTime = acprof->launchEntryTime;
        commTurnProf.waitMsgStartTime = acprof->commInitEndTime;
        commTurnProf.kfcAlgExeStartTime = commRcd->hccExecStartTime;
        commTurnProf.sendTaskStartTime = commRcd->sendTaskStartTime;
        commTurnProf.sendSqeFinishTime = commRcd->sendSqeFinishTime;
        commTurnProf.rtsqExeEndTime = acprof->receiveFinalizeTime;
        commTurnProf.serverEndTime = acprof->endTime;
        commTurnProf.dataLen = commRcd->dataLen;

        commTurnProf.deviceId = acprof->rankId;
        commTurnProf.streamId = streamId;
        commTurnProf.taskId = taskId;
        commTurnProf.commTurn = workRcdCnt;
        commTurnProf.currentTurn = j;
        commTurnProf.version = 0U; // default
        HCCL_INFO("Get one comm prof with details:[%s]",
            dfx::ProfilingExtendInfoHelper::AicpuKfcProfCommTurnToString(commTurnProf).c_str());
        CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(
            MSPROF_REPORT_AICPU_MC2_EXECUTE_COMM_TIME, GetCurCpuTimestamp(true), &commTurnProf, sizeof(commTurnProf)));
    }
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::ReportTaskInfo(hccl::Stream stream)
{
    if (!GetProfL1State()) {
        return HCCL_SUCCESS;
    }
    ProfCommInfo profInfo;
    auto streamId = stream.id();
    // 通信域中获取rankId, groupHashId等信息
    CHK_RET(GetProfInfoByStreamId(streamId, profInfo));
    hccl::HcclSqeContext *sqeContext = stream.GetSqeContextPtr();
    hccl::SqeRingBuffer *sqeContextBuffer = &(sqeContext->buffer);
    CHK_PTR_NULL(sqeContextBuffer);
    u32 startSqeIdx = GetStartReportSqeIdx(streamId);
    HCCL_INFO("[ReportTaskInfo] Rank:%u, stream:%u, sqeNum:%u, startSqeIdx:%u, curSqeTailIdx: %u", profInfo.rankId, streamId,
        sqeContextBuffer->tailSqeIdx - startSqeIdx, startSqeIdx, sqeContextBuffer->tailSqeIdx);
    MsprofAicpuHcclTaskInfo taskInfos[HCCLINFO_REPORT_BATCH_NUM] = {0};
    auto endIdx = static_cast<uint32_t>(sqeContextBuffer->tailSqeIdx);
    for (uint32_t idx = startSqeIdx, batchId = 0; idx < endIdx; ++idx) {
        // 获取SqeInfo
        SqeInfo sqeInfo{};
        SqeContextUtils::QuerySqeInfo(sqeContextBuffer->localBuff + idx * hccl::HCCL_SQE_SIZE,
            sqeContextBuffer->sqeType[idx], sqeContextBuffer->addInfo[idx], &sqeInfo);
        sqeInfo.remoteRank = sqeContextBuffer->dfxInfo[idx].remoteRank;
        // 转换为MsprofAicpuHcclTaskInfo
        auto& taskInfo = taskInfos[batchId++];
        dfx::ProfilingExtendInfoHelper::InitHcclInfo(taskInfo);
        CommInfo2HcclInfo(profInfo, taskInfo);
        dfx::ProfilingExtendInfoHelper::SqeInfo2MsprofAicpuMC2HcclInfo(sqeInfo, taskInfo);
        taskInfo.timeStamp = sqeContextBuffer->profTimestap[idx]; // 时间戳
        DumpHcclInfo(taskInfo, batchId, idx);
        // 上报信息
        if (batchId == HCCLINFO_REPORT_BATCH_NUM || idx == (endIdx - 1)) {
            CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(MSPROF_REPORT_AICPU_MC2_BATCH_HCCL_INFO,
                0, taskInfos, sizeof(MsprofAicpuHcclTaskInfo) * batchId));
            batchId = 0;
            memset_s(taskInfos, sizeof(taskInfos), 0, sizeof(taskInfos));
        }
    }
    CHK_RET(UpdateStartReportSqeIdx(streamId, sqeContextBuffer->tailSqeIdx));
    return HCCL_SUCCESS;
}

void ProfilingManager::DumpHcclInfo(const MsprofAicpuHcclTaskInfo& taskInfo, u32 batchId, u32 idx)
{
    HCCL_DEBUG("[ReportTaskInfo] batchId:%u, idx:%u, itemId:%lld, groupName:%lld, localRank:%u, remoteRank:%u, " \
            "rankSize:%u, timeStamp:%lld, srcAddr:%x, dstAddr:%x, dataSize:%lld, taskId:%u, streamId:%u, planeID:%u," \
            "opType:%u, dataType:%u, linkType:%u, transportType:%u, rdmaType:%u, role:%u",
            batchId, idx, taskInfo.itemId, taskInfo.groupName, taskInfo.localRank, taskInfo.remoteRank,taskInfo.rankSize,
            taskInfo.timeStamp, taskInfo.srcAddr, taskInfo.dstAddr, taskInfo.dataSize,taskInfo.taskId, taskInfo.streamId,
            taskInfo.planeID, taskInfo.opType, taskInfo.dataType, taskInfo.linkType, taskInfo.transportType,
            taskInfo.rdmaType, taskInfo.role);
}

void ProfilingManager::CommInfo2HcclInfo(const dfx::ProfCommInfo &profInfo, MsprofAicpuHcclTaskInfo &taskInfo)
{
    taskInfo.groupName = profInfo.groupNameHashId;
    taskInfo.localRank = profInfo.rankId;
    taskInfo.rankSize = profInfo.rankNum;
}

HcclResult ProfilingManager::ReportHcclOpInfo(MsprofAicpuHCCLOPInfo& hcclOpInfo, std::string &algTypeStr)
{
    if (!GetProfL0State()) {
        return HCCL_SUCCESS;
    }
    CHK_PTR_NULL(aicpu::GetTaskAndStreamId);
    uint64_t taskId = 0U;
    uint32_t streamId = 0;
    CHK_PRT_RET(aicpu::GetTaskAndStreamId(taskId, streamId) != aicpu::status_t::AICPU_ERROR_NONE,
        HCCL_ERROR("Failed to get task id and stream id."), HCCL_E_PARA);

    hcclOpInfo.algType = GetProfHashId(algTypeStr.c_str(), algTypeStr.length());
    hcclOpInfo.taskId = taskId;
    hcclOpInfo.streamId = streamId;
    HCCL_INFO("[ReportHcclOpInfo] relay:%u, retry:%u, dataType:%u, algType:%u, count:%llu, groupHashId:%llu",
        hcclOpInfo.relay, hcclOpInfo.retry, hcclOpInfo.dataType, hcclOpInfo.algType, hcclOpInfo.count,
        hcclOpInfo.groupName);
    CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(MSPROF_REPORT_AICPU_HCCL_OP_INFO,
        ProfGetCurCpuTimestamp(), &hcclOpInfo, sizeof(MsprofAicpuHCCLOPInfo)));
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::ReportMainStreamTask(hccl::Stream& stream, uint16_t taskId, uint16_t type)
{
    if (!GetProfL0State()) {
        return HCCL_SUCCESS;
    }
    CHK_PTR_NULL(aicpu::GetTaskAndStreamId);
    uint64_t aicpuKernelTaskId = 0U;
    uint32_t aicpuKernelStreamId = 0;
    CHK_PRT_RET(aicpu::GetTaskAndStreamId(aicpuKernelTaskId, aicpuKernelStreamId) != aicpu::status_t::AICPU_ERROR_NONE,
        HCCL_ERROR("Failed to get task id and stream id."), HCCL_E_PARA);
    MsprofAicpuHcclMainStreamTask flagtask{};
    flagtask.streamId = stream.id();
    flagtask.taskId = taskId;
    flagtask.type = type;
    flagtask.aicpuStreamId = aicpuKernelStreamId;
    flagtask.aicpuTaskId = aicpuKernelTaskId;
    HCCL_INFO("[ReportMainStreamTask] streamId:%u, taskId:%u, type:%u", flagtask.streamId, flagtask.taskId, flagtask.type);

    CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(MSPROF_REPORT_AICPU_HCCL_FLAG_TASK,
        ProfGetCurCpuTimestamp(), &flagtask, sizeof(MsprofAicpuHcclMainStreamTask)));

    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::ReportFilpTask(s32 streamId, uint16_t taskId, uint32_t flipNum)
{
    if (!GetProfL0State()) {
        return HCCL_SUCCESS;
    }
    MsporfAicpuFlipTask flipTaskInfo{};
    flipTaskInfo.streamId = streamId;
    flipTaskInfo.taskId = taskId;
    flipTaskInfo.flipNum = flipNum;
    HCCL_INFO("[ReportFilpTask] streamId:%u, taskId:%u, filpNum:%u", flipTaskInfo.streamId, flipTaskInfo.taskId,
        flipTaskInfo.flipNum);
    CHK_PRT(dfx::ProfilingManager::CallMsprofReportAdditionInfo(MSPROF_REPORT_AICPU_FILP_TASK,
        ProfGetCurCpuTimestamp(), &flipTaskInfo, sizeof(MsporfAicpuFlipTask)));

    return HCCL_SUCCESS;
}

uint64_t ProfilingManager::GetProfHashId(const char *name, uint32_t len)
{
    if (name == nullptr || len == 0) {
        HCCL_WARNING("HashData is empty.");
        return INVALID_U64;
    }
    return AdprofGetHashId(name, len);
}

uint32_t ProfilingManager::GetStartReportSqeIdx(s32 streamId)
{
    std::unique_lock<std::mutex> lock(startReportSqeIdxMutex_);
    u32 lastSqeTailIdx = 0;
    auto iter = streamToSqeIdxMap_.find(streamId);
    if (iter == streamToSqeIdxMap_.end()) {
        HCCL_INFO("[GetProfInfoByStreamId]streamId:%d is not find", streamId);
        streamToSqeIdxMap_.insert({streamId, 0});
    } else {
        lastSqeTailIdx = iter->second;
    }
    return lastSqeTailIdx;
}

HcclResult ProfilingManager::UpdateStartReportSqeIdx(s32 streamId, u32 newSqeTailIdx)
{
    std::unique_lock<std::mutex> lock(startReportSqeIdxMutex_);
    auto iter = streamToSqeIdxMap_.find(streamId);
    if (iter == streamToSqeIdxMap_.end()) {
        streamToSqeIdxMap_.insert({ streamId, newSqeTailIdx });
        HCCL_INFO("[UpdateStartReportSqeIdx]streamId:%d is not find, newSqeTailIdx:%u", streamId, newSqeTailIdx);
    } else {
        // 到2048时，更新成0;
        newSqeTailIdx = (iter->second == hccl::HCCL_SQE_MAX_CNT) ? 0 : newSqeTailIdx;
        HCCL_INFO("[UpdateStartReportSqeIdx] streamId:%d, lastSqeTailIdx:%u, newSqeTailIdx:%u", streamId,
            iter->second, newSqeTailIdx);
        iter->second = newSqeTailIdx;
    }
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::GetProfInfoByStreamId(s32 streamId, ProfCommInfo& profInfo)
{
    std::string tag = "unkown";
    std::unique_lock<std::mutex> lock(streamMutex_);
    auto iter = streamToTagMap_.find(streamId);
    if (iter == streamToTagMap_.end()) {
        HCCL_INFO("[GetProfInfoByStreamId]streamId:%d is not find", streamId);
    } else {
        tag = iter->second;
        auto opInfoIter = tagOpInfoMap_.find(tag);
        if (opInfoIter == tagOpInfoMap_.end()) {
            HCCL_INFO("[GetProfInfoByStreamId]streamId:%s is not find", tag.c_str());
        } else {
            profInfo = opInfoIter->second;
        }
    }
    return HCCL_SUCCESS;
}

HcclResult ProfilingManager::AddProfInfoByStreamId(s32 streamId, const std::string &tag, const ProfCommInfo& profInfo)
{
    std::unique_lock<std::mutex> lock(streamMutex_);
    auto tagMapIter = streamToTagMap_.find(streamId);
    if (tagMapIter == streamToTagMap_.end()) {
        streamToTagMap_.insert({streamId, tag});
    } else {
        tagMapIter->second = tag;
        // streamId之前存在说明流被销毁了，被其它通信域复用了, 上报一条信息告知profiling taskid发生翻转
        CHK_RET(ReportFilpTask(streamId, UINT16_MAX, UINT16_MAX));
        HCCL_INFO("[AddProfInfoByStreamId] streamId:%d content:%s update", streamId, tag.c_str());
    }

    auto opInfoMapIter = tagOpInfoMap_.find(tag);
    if (opInfoMapIter == tagOpInfoMap_.end()) {
        tagOpInfoMap_.insert(std::make_pair(tag, profInfo));
    } else {
        opInfoMapIter->second = profInfo;
    }
    return HCCL_SUCCESS;
}

}  // namespace dfx